#include "webclient.hpp"
#include <boost/beast/core/buffers_to_string.hpp>
#include <iostream>
#include <sstream>

// 构造函数
WebClient::WebClient()
    : ctx(ssl::context::tlsv12_client)
    , resolver(ioc)
    , isSecure(false)
{
    new (&ws) websocket::stream<tcp::socket>(ioc);
}

WebClient::WebClient(std::string host, std::string port, std::string path, bool isSecure)
    : host(host)
    , port(port)
    , path(path)
    , ctx(ssl::context::tlsv12_client)
    , resolver(ioc)
    , isSecure(isSecure)
{
    if (isSecure) {
        new (&wss) websocket::stream<beast::ssl_stream<tcp::socket>>(ioc, ctx);
    } else {
        new (&ws) websocket::stream<tcp::socket>(ioc);
    }
}

WebClient::WebClient(std::string url)
    : url(url)
    , ctx(ssl::context::tlsv12_client)
    , resolver(ioc)
{
    parseUrl(url);
    if (isSecure) {
        new (&wss) websocket::stream<beast::ssl_stream<tcp::socket>>(ioc, ctx);
    } else {
        new (&ws) websocket::stream<tcp::socket>(ioc);
    }
}

// 析构函数
WebClient::~WebClient()
{
    try {
        close();
    } catch (...) {
        // 忽略析构时的异常
    }
    if (isSecure) {
        wss.websocket::stream<beast::ssl_stream<tcp::socket>>::~stream();
    } else {
        ws.websocket::stream<tcp::socket>::~stream();
    }
}

// 配置函数
void WebClient::config(std::string host, std::string port, std::string path, bool isSecure)
{
    this->host = host;
    this->port = port;
    this->path = path;
    this->isSecure = isSecure;
    if (isSecure) {
        if (!this->isSecure) {
            ws.websocket::stream<tcp::socket>::~stream();
        }
        new (&wss) websocket::stream<beast::ssl_stream<tcp::socket>>(ioc, ctx);
    } else {
        if (this->isSecure) {
            wss.websocket::stream<beast::ssl_stream<tcp::socket>>::~stream();
        }
        new (&ws) websocket::stream<tcp::socket>(ioc);
    }
}

void WebClient::config(std::string url)
{
    this->url = url;
    parseUrl(url);
    if (isSecure) {
        if (!this->isSecure) {
            ws.websocket::stream<tcp::socket>::~stream();
        }
        new (&wss) websocket::stream<beast::ssl_stream<tcp::socket>>(ioc, ctx);
    } else {
        if (this->isSecure) {
            wss.websocket::stream<beast::ssl_stream<tcp::socket>>::~stream();
        }
        new (&ws) websocket::stream<tcp::socket>(ioc);
    }
}

// 连接函数
int WebClient::connect()
{
    if (connected) {
        return -1;
    }

    beast::error_code ec;

    // 解析主机名和端口
    auto const results = resolver.resolve(host, port, ec);
    if (ec) {
        std::cerr << "Resolve failed: " << ec.message() << std::endl;
        return -2;
    }

    if (isSecure) {
        // 连接到服务器
        net::connect(wss.next_layer().next_layer(), results.begin(), results.end(), ec);
        if (ec) {
            std::cerr << "Connect failed: " << ec.message() << std::endl;
            return -3;
        }

        if (!SSL_set_tlsext_host_name(wss.next_layer().native_handle(), host.c_str())) {
            std::cerr << "Error: SNI setting failed" << std::endl;
            return -4;
        }

        // TLS 握手
        wss.next_layer().handshake(ssl::stream_base::client, ec);
        if (ec) {
            std::cerr << "TLS Handshake failed: " << ec.message() << std::endl;
            return -5;
        }

        // 执行 WebSocket 握手
        wss.handshake(host, path, ec);
        if (ec) {
            std::cerr << "Handshake failed: " << ec.message() << std::endl;
            return -6;
        }
    } else {
        // 连接到服务器
        net::connect(ws.next_layer(), results.begin(), results.end(), ec);
        if (ec) {
            std::cerr << "Connect failed: " << ec.message() << std::endl;
            return -3;
        }

        // 执行 WebSocket 握手
        ws.handshake(host, path, ec);
        if (ec) {
            std::cerr << "Handshake failed: " << ec.message() << std::endl;
            return -4;
        }
    }

    connected = true;
    return 0;
}

// 发送消息函数
int WebClient::send(const std::string& message)
{
    if (!connected) {
        return -1;
    }

    beast::error_code ec;

    if (isSecure) {
        wss.write(net::buffer(message), ec);
    } else {
        ws.write(net::buffer(message), ec);
    }
    if (ec) {
        std::cerr << "Send message failed: " << ec.message() << std::endl;
        return -2;
    }

    return 0;
}

// 接收消息函数
int WebClient::receive(std::string& received)
{
    if (!connected) {
        return -1;
    }

    beast::error_code ec;
    beast::flat_buffer buffer;
    if (isSecure) {
        wss.read(buffer, ec);
    } else {
        ws.read(buffer, ec);
    }
    if (ec) {
        std::cerr << "Receive message failed: " << ec.message() << std::endl;
        return -2;
    }

    received = beast::buffers_to_string(buffer.data());
    return 0;
}

// 关闭连接函数
void WebClient::close()
{
    if (!connected) {
        return;
    }

    beast::error_code ec;
    if (isSecure) {
        wss.close(websocket::close_code::normal, ec);
    } else {
        ws.close(websocket::close_code::normal, ec);
    }
    if (ec) {
        std::cerr << "Close failed: " << ec.message() << std::endl;
    }
    connected = false;
}

// 解析 URL 函数
void WebClient::parseUrl(const std::string& url)
{
    std::string protocol;
    std::istringstream iss(url);
    std::getline(iss, protocol, ':');
    if (protocol == "ws") {
        isSecure = false;
        port = "80"; // 默认端口
    } else if (protocol == "wss") {
        isSecure = true;
        port = "443"; // 默认 WSS 端口
    } else {
        throw std::invalid_argument("Invalid protocol. Expected 'ws' or 'wss'.");
    }
    iss.ignore(2); // 跳过 "://"

    std::string hostPort;
    std::getline(iss, hostPort, '/');
    size_t colonPos = hostPort.find(':');
    if (colonPos != std::string::npos) {
        this->host = hostPort.substr(0, colonPos);
        this->port = hostPort.substr(colonPos + 1);
    } else {
        this->host = hostPort;
    }

    this->path = "/" + iss.str().substr(iss.tellg());
}